from __future__ import annotations
from PIL import Image, ImageDraw, ImageFont
from ...dto.image import ImageGenerationRequest
from typing import List, Tuple
from threading import Thread, Event
from queue import Queue
import asyncio
from .simulate import concurrent_simulate


class GeneratePorter:
    """
    新起一个线程来处理并发的图像生成请求, 并按照顺序返回结果
    利用asyncio.Future来模拟异步图像生成
    """

    _instance: GeneratePorter = None
    __slots__ = ("_queue", "_lock", "_thread", "_stop_event")

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(GeneratePorter, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if hasattr(self, "_queue"):
            return
        self._queue: Queue[Tuple[asyncio.Future, ImageGenerationRequest]] = Queue()
        self._lock: asyncio.Lock = asyncio.Lock()
        self._thread: Thread = Thread(target=self.run, daemon=True)
        self._thread.start()

    def run(self):
        self._stop_event = Event()
        while not self._stop_event.is_set():
            future, request = self._queue.get()
            try:
                match request.model: # 根据模型类型分发任务(可以被模型实例决定处理批还是单个请求, 内部set_result即可)
                    case "dall-e-3":
                        batch_dall_e3(future, request)
                    case "simulate": # 示例攒批
                        concurrent_simulate(future, request)
                    case _:
                        raise ValueError(f"Unsupported model: {request.model}")
            except Exception as e:
                future.set_exception(e)
            finally:
                self._queue.task_done()

    def stop(self):
        self._stop_event.set()
        self._thread.join()
        self.__class__._instance = None

    def put(self, future: asyncio.Future, request: ImageGenerationRequest):
        self._queue.put((future, request))


def generate_image_entrypoint(
    request: ImageGenerationRequest,
) -> List[asyncio.Future[Image.Image]]:
    """
    模拟一个图像生成函数。
    在实际应用中，您将在这里集成您的文生图模型。
    """
    futures = []
    generator = GeneratePorter()  # 实例化生成器(单例模式)
    for idx in range(request.n):
        future = asyncio.Future()
        generator.put(future, request)  # 放入队列等待生成
        futures.append(future)
    return futures

def batch_dall_e3(
    future: asyncio.Future,
    request: ImageGenerationRequest
):
    width, height = map(int, request.size.split("x"))
    prompt = request.prompt
    img = Image.new("RGB", (width, height), color="white")
    d = ImageDraw.Draw(img)
    font = ImageFont.load_default()
    d.text((10, 10), f"Prompt: {prompt[:50]}...", fill=(0, 0, 0), font=font)
    d.text(
        (10, 50),
        f"Size: {request.size}, Quality: {request.quality}",
        fill=(0, 0, 0),
        font=font,
    )
    future.set_result(img)  # 设置生成的图像结果